# -*- coding: utf-8 -*-
"""
Created on Wed Sep 23 10:41:00 2020

@author: Robin
"""
from .utils import process_raw_data
from .IOFile import readjson
#from DataIO import *

import pandas as pd
import re

readpath = 'data'

'''
根据传入参数来获取资产数据
''' 
def cal_stockdf(stock_codes,begdate,enddate,
                 period=1,
                 include_st = False,
                 include_suspend = False,
                 include_new_stock = False,
                 ipo_days = 60):         
    # 读取数据
    stock_file_name = 'stockPrice_' + re.sub('[.]','',stock_codes)
    prices = readjson(stock_file_name,readpath)
    prices_use=prices.loc[begdate:enddate]
    
    return prices_use
    '''
    cal_stockdf的返回值prices_useSH是一个DataFrame，索引是时间，列名是股票代码，内容是收盘价
                000008.SZ  000009.SZ   000012.SZ  ...  603888.SH  603939.SH  603983.SH
    2020-09-10  67.673045  60.038643  140.051163  ...  61.867993  279.44919  65.306056
    2020-09-11  67.448962  61.548490  143.532358  ...  62.393411  288.53064  66.541657
    2020-09-14  67.448962  63.680040  144.067926  ...  60.107842  287.34861  66.029334
    '''


'''
根据传入参数来获取因子的数据
'''      
def cal_inddf(stock_codes,ind_codes,begdate='1900-01-01',enddate='2099-12-31',
             period=1):  
        
    #file_weight_info输入行业信息表的名称
    
    #begdate和enddate输入时为字符串 格式为"2017-6-19"这样
#    begdate=pd.to_datetime(begdate)
#    enddate=pd.to_datetime(enddate)
    
    # 只提取一个因子
    if isinstance(ind_codes,str):
        # 读取数据
        indi = readjson(ind_codes,readpath)
   
        #选取开始日至结束日之间的因子数据，频率为period（1,2,3...）
        indi_use=indi[begdate:enddate:period]
        # 转为双索引
        
        time_stock_Index = pd.MultiIndex.from_product(
                [indi_use.index,indi_use.columns],
                 names = ['date','asset'])
        factors = pd.DataFrame({ind_codes:indi_use.values.reshape(indi_use.size)},index = time_stock_Index)
        
    # 提取多个因子
    elif isinstance(ind_codes,list):     
        factors = {}
        for ind_code in ind_codes:
            # 读取数据
            indi = readjson(ind_code,readpath)
            indi_use=indi[begdate:enddate:period]
            # 转为双索引
            
            time_stock_Index = pd.MultiIndex.from_product(
                    [indi_use.index,indi_use.columns],
                     names = ['date','asset'])
            factors[ind_code] = indi_use.values.reshape(indi_use.size)
        
        # 考虑到不同因子可能存储的频率不同，所以需要改用merge等来实现索引的统一          
        factors = pd.DataFrame(factors,index = time_stock_Index)
    
    
    #调用utils.py里的process_row_data2函数进行数据标准化
    factors=process_raw_data(factors)
    
    return factors
    '''
    cal_inddf的返回值factors是一个双索引的DataFrame,索引是（日期，股票代码）
    列名是因子名称，内容是因子值
                                indi1      indi2
    date       asset                  
    2020-09-10 000008.SZ      1.836478     675.8
               000009.SZ      0.501466     13.97
               000012.SZ      1.375007     ....
               
    2020-09-22 603328.SH     -0.919232     ....
               603338.SH     -1.057981     ....
               603355.SH      1.652590     ....
    '''



'''导入指数（中证500）的收盘价数据'''
def cal_indexdf(index_name,begdate,enddate,
                 period=1):         
    # 读取数据
    stock_file_name = 'index_close_' + re.sub('[.]','',index_name)
    #index_name输入指数的名称 在这里是000905.SH_in_day（因为文件名是这样命名的）
    prices = readjson(stock_file_name,readpath)
    prices_use=prices.loc[begdate:enddate]
    
    return prices_use
    '''
    cal_indexdf的返回值prices_use是一个DataFrame，内容是指数的收盘价
                    close
    2020-09-10  6235.5687
    2020-09-11  6311.9670
    2020-09-14  6351.2422
    '''


'''导入股票市值的数据'''
def cal_market_value_df(stock_codes,begdate,enddate,
                 period=1,
                 include_st = False,
                 include_suspend = False,
                 include_new_stock = False,
                 ipo_days = 60):         
    # 读取数据
    stock_file_name = 'market_value_' + re.sub('[.]','',stock_codes)
    prices = readjson(stock_file_name,readpath)
    prices_use=prices.loc[begdate:enddate]
    
    return prices_use
    '''
    cal_market_value_df的返回值是一个DataFrame,索引是日期，列名是股票代码，内容是股票市值
                   000008.SZ     000009.SZ  ...     603939.SH    603983.SH
    2020-09-11  8.370194e+09  1.787395e+10  ...  5.314932e+10  26562240000
    2020-09-14  8.370194e+09  1.849296e+10  ...  5.293158e+10  26357730000
    2020-09-15  8.342386e+09  1.859613e+10  ...  5.441326e+10  26987300000
    '''


'''
行业数据
'''
def cal_groupby(stock_codes,
                begdate='1900-01-01',
                enddate='2099-12-31',   
                period = 1,
                industry_type='sw_1'):
    '''
    groupby : pd.Series - MultiIndex or dict
                Either A MultiIndex Series indexed by date and asset,
                containing the period wise group codes for each asset, or
                a dict of asset to group mappings. If a dict is passed,
                it is assumed that group mappings are unchanged for the
                entire time period of the passed factor data.
    ''' 
    filename = 'index_components_000905SH'
    groupby = readjson(filename,readpath)
    groupby.set_index(pd.MultiIndex.from_frame(groupby[['date','asset']]),inplace=True)
    
    return groupby['industry']


'''
个股权重数据
'''
def cal_weightdf(stock_codes,
                 base_code='000905.SH',
                 begdate='1900-01-01',
                 enddate='2099-12-31',
                 period = 1):
    
    filename = 'index_components_000905SH'
    weight_data = readjson(filename,readpath)
    weight_data.set_index(pd.MultiIndex.from_frame(weight_data[['date','asset']]),inplace=True)
    return weight_data['i_weight']  
    '''
    cal_weightdf的返回值是一个双索引的DataFrame，索引是（日期，股票代码），内容是
    个股在中证500中的权重数据（官方数据）
    date        asset    
    2007-01-31  000006.SZ    0.470
                000016.SZ    0.240
                000028.SZ    0.140
    2020-08-31  603328.SH    0.070
                603338.SH    0.367
                603355.SH    0.044
    '''



def cal_inddf_with_indu(ind_df,industry_type='sw_1',industry_level = 1):
    factors = ind_df.copy(deep = True)
    # 加入行业分类信息
    if industry_type != None:
        # 按照行业分类方法来添加行业信息
        file_weight_info = 'index_components_000905SH'
        weight_info = readjson(file_weight_info,readpath)
        weight_info=weight_info.set_index([weight_info['date'], weight_info['asset']],drop=False)
        '''
        weight_info是一个双索引的DataFrame，储存个股的行业信息和权重信息，格式如下：
                                    'industry'    'weight'
        2020-01-03 '000008.sz'         行业1      0.840348
        2020-01-03 '000009.sz'         行业2      0.341126
                     ....               ...         ....
        2020-01-10 '000008.sz'         行业1      0.707499
        2020-01-10 '000009.sz'         行业2      0.312943
                     ....               ...         .... 
        '''

#此处待修改。现在导入的weight_info索引和factors不一样    
    #weight_info_use=weight_info[begdate:enddate:period]
#    #现在weight_info_use的索引（双索引，日期和股票代码）和factors的索引完全一样
#    factors['industry']=weight_info['industry']
#    factors['weight']=weight_info['weight']    
                 
##测试数据。暂时没有个股的行业数据，随便编造了一个df替代
#weights=list(np.random.rand(19000))
#industry_list=[]
#for i in range(19000):
#    industry_list.append('测试行业')
#
#df_industry_weight=pd.DataFrame()
#df_industry_weight['industry']=industry_list
#df_industry_weight['weight']=weights
#df_industry_weight.index=factors.index
    return factors




